import os
import joblib
import torch
import numpy as np
from torchvision import transforms, models
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm
from PIL import Image
from sklearn.svm import SVC

def extract_features(model, dataloader, device):
    features, labels = [], []
    model.eval()
    with torch.no_grad():
        for imgs, lbls in tqdm(dataloader, desc="Extracting features", unit="batch"):
            imgs = imgs.to(device)
            feat = model(imgs)
            features.append(feat.cpu().numpy())
            labels.append(lbls.numpy())
    return np.vstack(features), np.hstack(labels)


class DogDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.images, self.labels = [], []
        self.classes = sorted(os.listdir(root_dir))
        self.class_to_idx = {cls: i for i, cls in enumerate(self.classes)}

        for cls in self.classes:
            cls_path = os.path.join(root_dir, cls)
            for img_name in os.listdir(cls_path):
                if img_name.lower().endswith((".jpg", ".jpeg", ".png")):
                    self.images.append(os.path.join(cls_path, img_name))
                    self.labels.append(self.class_to_idx[cls])

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        img = Image.open(self.images[idx]).convert("RGB")
        label = self.labels[idx]
        if self.transform:
            img = self.transform(img)
        return img, label


def train_svm_model(train_path, test_path):
    print(f"📁 训练路径: {train_path}")
    print(f"📁 测试路径: {test_path}")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("⚙️ 使用设备:", device)

    transform = transforms.Compose([
        transforms.Resize(299),
        transforms.CenterCrop(299),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])

    # 1️⃣ 加载数据集
    train_dataset = DogDataset(train_path, transform)
    test_dataset = DogDataset(test_path, transform)
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=False, num_workers=0)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=0)

    # 2️⃣ 加载模型（去掉最后分类层）
    model = models.inception_v3(pretrained=True)
    model.fc = torch.nn.Identity()
    model.to(device)

    # 3️⃣ 提取特征
    print("🔍 正在提取训练集特征...")
    X_train, y_train = extract_features(model, train_loader, device)
    print("🔍 正在提取测试集特征...")
    X_test, y_test = extract_features(model, test_loader, device)

    # 4️⃣ SVM 训练
    print("🚀 开始训练 SVM 模型...")
    svm = SVC(kernel='rbf', decision_function_shape='ovo')
    svm.fit(X_train, y_train)
    print("✅ 训练完成")

    # 5️⃣ 测试准确率
    y_pred = svm.predict(X_test)
    accuracy = (y_pred == y_test).mean()
    print(f"🎯 测试准确率: {accuracy * 100:.2f}%")

    # 6️⃣ 保存模型
    os.makedirs("models", exist_ok=True)
    model_path = f"models/svm_model_{int(accuracy*100)}.pkl"
    joblib.dump(svm, model_path)
    print(f"💾 模型已保存: {model_path}")

    return {"accuracy": round(accuracy * 100, 2), "modelPath": model_path}
